# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/18 19:00 
@Author  : ZhangShenao 
@File    : 3.使用CallbackHandler实现回调机制.py 
@Desc    : 使用LangChain内置的CallbackHandler,来实现回调机制

LangChain通过回调机制提供了一种灵活的方式
来监控和操作应用程序的不同阶段

"""
import os

import dotenv
from langchain_core.callbacks import FileCallbackHandler, StdOutCallbackHandler
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from loguru import logger

# 加载环境变量
dotenv.load_dotenv()

# 初始化日志参数
log_file = './logs/out.log'
logger.add(log_file, colorize=True, enqueue=True)

# 创建CallbackHandler监听器
# FileCallbackHandler:将回调信息写入日志文件
# StdOutCallbackHandler:将回调信息打印到标准输出(控制台)
file_callback_handler = FileCallbackHandler(filename=log_file)
std_out_callback_handler = StdOutCallbackHandler()

# 创建LLM
llm = ChatOpenAI(model_name='gpt-3.5-turbo', openai_api_base=os.getenv('OPENAI_API_BASE'))

# 编排Prompt
prompt = ChatPromptTemplate.from_template('1 + {num} = ?')

# 创建OutputParser
parser = StrOutputParser()

# 创建Chain
chain = prompt | llm | parser

# 执行Chain
# 通过config配置参数传入CallbackHandler
output = chain.invoke(input={'num': 10},
                      config={'callbacks': [file_callback_handler, std_out_callback_handler]}
                      )

# 打印结果
print(output)
